package com.team1.system.model;
import java.util.*;

//信息检索:使用了向量空间模型
public class InformationRetrieval {
    private Map<String, Map<String, Integer>> index;
    private Map<String, Integer> documentLengths;
    private Set<String> stopWords;

    public InformationRetrieval() {
        index = new HashMap<>();//映射来存储每个文档中每个词的出现次数
        documentLengths = new HashMap<>();// 存储每个文档的长度
        stopWords = new HashSet<>();// 存储停用词。
    }

    //添加一个文档，documentId:文档的唯一标识符，documentText:文档的文本内容
    public void addDocument(String documentId, String documentText) {
        Map<String, Integer> termFrequencies = new HashMap<>();
        String[] terms = documentText.split("\\s+");
        int documentLength = 0;
        for (String term : terms) {
            term = term.toLowerCase();
            if (!stopWords.contains(term)) {
                termFrequencies.put(term, termFrequencies.getOrDefault(term, 0) + 1);
                documentLength++;
            }
        }
        index.put(documentId, termFrequencies);
        documentLengths.put(documentId, documentLength);
    }
    //添加一个停用词，这个停用词将在搜索时被忽略。
    public void addStopWord(String stopWord) {
        stopWords.add(stopWord.toLowerCase());
    }

    //对系统中的文档进行搜索，query是查询字符串。返回一个按照相关度排序的文档列表
    public List<String> search(String query) {
        Map<String, Double> scores = new HashMap<>();
        String[] terms = query.split("\\s+");
        for (String term : terms) {
            term = term.toLowerCase();
            if (!stopWords.contains(term)) {
                double idf = Math.log((double) index.size() / (double) getDocumentFrequency(term));
                for (String documentId : index.keySet()) {
                    double tf = getTermFrequency(term, documentId);
                    double score = tf * idf;
                    scores.put(documentId, scores.getOrDefault(documentId, 0.0) + score);
                }
            }
        }
        List<String> results = new ArrayList<>(scores.keySet());
        results.sort((a, b) -> Double.compare(scores.get(b), scores.get(a)));
        return results;
    }

    private int getDocumentFrequency(String term) {
        int count = 0;
        for (Map<String, Integer> termFrequencies : index.values()) {
            if (termFrequencies.containsKey(term)) {
                count++;
            }
        }
        return count;
    }

    private int getTermFrequency(String term, String documentId) {
        Map<String, Integer> termFrequencies = index.get(documentId);
        if (termFrequencies == null) {
            return 0;
        }
        return termFrequencies.getOrDefault(term, 0);
    }

    public static void main(String[] args) {
        InformationRetrieval irs = new InformationRetrieval();
        irs.addStopWord("the");
        irs.addStopWord("and");
        irs.addStopWord("of");
        irs.addDocument("doc1", "The quick brown fox jumps over the lazy dog");
        irs.addDocument("doc2", "The quick brown fox jumps over the quick dog");
        irs.addDocument("doc3", "Brown fox brown dog");
        List<String> results = irs.search("quick brown fox");
        for (String result : results) {
            System.out.println(result);
        }
    }
}